import sys
import os
import torch
from transformers import GPT2LMHeadModel, GPT2Config
from transformers import BertTokenizerFast

# 获取项目根目录 (不这么搞的话，导包不了  from gpt2_chatbot import *****)
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
print(f'project_root-->\n{project_root}')
sys.path.append(project_root)
from gpt2_chatbot.config import ParameterConfig





# 1 初始化 config
params = ParameterConfig()


# 2 初始化 tokenizer
tokenizer = BertTokenizerFast(params.vocab_path,
                                sep_token="[SEP]",
                                pad_token="[PAD]",
                                cls_token="[CLS]")


# 3 初始化 model
if params.pretrained_model:
    # 加载预训练模型
    model = GPT2LMHeadModel.from_pretrained(params.pretrained_model)
else:
    # 初始化模型
    model_config = GPT2Config.from_json_file(params.config_json)
    print(f'model的配置信息--->\n{model_config}')
    model = GPT2LMHeadModel(config=model_config)
    print(f'模型的结构--->\n{model}')




# 4 计算模型参数数量
num_parameters = 0
parameters = model.parameters()
for idx, parameter in enumerate(parameters):
    # print(f'idx-->{idx} num--->{parameter.numel()}') # 单独打印每层的参数
    num_parameters += parameter.numel()
print(f'模型参数总量--->{num_parameters}')


# 5 打印模型的数据类型
print(f'模型参数类型--->{model.dtype}')


# 6 打印模型中所有子模块的类型
print(f'模型子模块类型--->')
for name, module in model.named_modules():
    print(f"{name:40} --> {type(module)}")
    # if 'attn' in name or 'mlp' in name:
    #     print(f"{name:40} --> {type(module)}")


# 7 打印其中一个层
print(f'model.transformer.h[0]---->\n{model.transformer.h[0]}')


# 8 模型大小
total_params = sum(p.numel() for p in model.parameters())
size_MB = total_params * 4 / 1024 / 1024
print(f"参数总数----->: {total_params}")
print(f"模型大小（float32）----->: {size_MB:.2f} MB")


# 8 保存模型 只保存参数
torch.save(model.state_dict(), "gpt2_model.pth")

